from langchain_community.llms import Ollama
from langchain_community.embeddings import OllamaEmbeddings

from promptAdmin.prompt_app.LLM.base import LlmBase


class OllamaModel(LlmBase):

    def get_model(self):
        llm = Ollama(model="llama2-chinese")
        return llm

    def get_embedding(self):
        embeddings = OllamaEmbeddings(model="llama2-chinese")
        return embeddings